#include "commons.h"
#include "tasks.h"
#include "zxLib.h"

float odorLength = 1.0;
unsigned int zxAutoStop = 0;
unsigned int zxAutoSwitch = 0;
unsigned int laserTimer = 0;
unsigned int laserOnTime = 20;
unsigned int laserOffTime = 30;
unsigned int laserTimerOn = 0;
unsigned int licking = 0;

unsigned int laserTrialType = laserOtherTrial;
//unsigned int ZJAllSessionLaserOn = 0;

void InitZXTMR(void) {
    TMR1 = 0;
    PR1 = TMR_RATE / 200;
    IFS0bits.T1IF = 0;
    IEC0bits.T1IE = 1;
    T1CON = 0x8020;
    ConfigIntTimer2(T1_INT_PRIOR_5 & T1_INT_ON);
}

void __attribute__((__interrupt__, no_auto_psv)) _T1Interrupt(void) {
    IFS0bits.T1IF = 0;
    laserTimer += 5;

    if (laserTimerOn && (laserTimer % (laserOnTime + laserOffTime) < laserOnTime)) {
        Out4 = 1;
    } else {
        Out4 = 0;
    }

    if (LickOnPin && !licking) {
        licking = 1;
        localSendOnce(SpLick, 1);
        Valve_ON(water_bitter, period * 2);
    } else if (licking && !LickOnPin) {
        licking = 0;
        Valve_OFF(water_bitter);
    }
}

void setLaser(void) {
    int n = getFuncNumber(1, "DC=1 20Hz=2");
    switch (n) {
        case 1:
            laserOnTime = 65535;
            laserOffTime = 0;
            break;
        case 2:
            laserOnTime = 20;
            laserOffTime = 30;
            break;
        default:
            laserOnTime = 65535;
            laserOffTime = 0;
    }
}

void protectedSerialSend(int type, int value) {
    IEC0bits.T1IE = 0;
    IEC0bits.T3IE = 0;
    DelayNmSecNew(4);
    localSendOnce(type, value);
    DelayNmSecNew(4);
    IEC0bits.T1IE = 1;
    IEC0bits.T3IE = 1;
}

void turnOnLaser() {
    LCDWriteChar('L', 4, 1);
    laserTimerOn = 1;
}

void turnOffLaser() {
    LCDWriteChar('.', 4, 1);
    laserTimerOn = 0;
}

void assertLaser(int type, int step, int currentTrial) {
    if (currentTrial) {
        switch (type) {
            case laserDuringBeginningToOneSecInITI:
                if (step == atFirstOdorBeginning) {
                    turnOnLaser();
                } else if (step == atITIOneSecIn) {
                    turnOffLaser();
                }
                break;
            case laserDuringDelay:
                if (step == atDelayOneSecIn) {
                    turnOnLaser();
                } else if (step == atDelayLastSecBegin) {
                    turnOffLaser();
                }
                break;
            case laserDuringOdor:
                if (step == atFirstOdorBeginning || step == atSecondOdorBeginning) {
                    turnOnLaser();
                } else if (step == atFirstOdorEnd || step == atSecondOdorEnd) {
                    turnOffLaser();
                }
        }
    }
}

void zxDNMSShaping(void) {
    wait_ms(1000);

    int currentTrial = 0;
    int currentSession = 0;
    int totalSession = 50;
    int totalTrial = 20;

    while (currentSession < totalSession) {
        currentSession++;
        protectedSerialSend(SpSess, 1);

        home_clr();
        puts_lcd((unsigned char*) &text7[0], sizeof (text7));
        line_2();
        puts_lcd((unsigned char*) &text8[0], sizeof (text8));

        LCDWriteNumber(currentSession, 2, 2, 2);
        currentTrial = 0;
        hit = 0;
        miss = 0;
        falseAlarm = 0;
        correctRejection = 0;

        correctRatio = 0;
        int currentTrial = 0;

        while (currentTrial < totalTrial) {
            currentTrial++;

            int r = rand() % 2;
            int firstOdor = r ? odor_B : odor_A;
            int secondOdor = r ? odor_A : odor_B;
            LCD_set_xy(1, 1);
            lcd_data(firstOdor + 0x3F);
            lcd_data(secondOdor + 0x3F);

            int laserOnTrial = (laserTrialType == laserNoTrial) ? 0 : (laserTrialType == laserEveryTrial) ? 1 : (currentTrial % 2);
            //            LCDWriteNumber(laserOnTrial, 1, 3, 1);
            zxDNMSLaserTrial(laserDuringBeginningToOneSecInITI, firstOdor, 1.0, 5.0, secondOdor, 0.6, 10, 1.0, laserOnTrial);
        }
        protectedSerialSend(SpSess, 0);
    }
}
//
//void setOdorLength(float in) {
//    odorLength = in;
//}

void shuffleArray(unsigned int orgArray[], unsigned int arraySize) {
    if (arraySize == 0 || arraySize == 1)
        return;

    int iter;
    for (iter = 0; iter < arraySize; iter++) {
        orgArray[iter] = iter;
    }
    int index, temp;
    for (iter = arraySize - 1; iter > 0; iter--) {
        index = rand() % (iter + 1);
        temp = orgArray[index];
        orgArray[index] = orgArray[iter];
        orgArray[iter] = temp;
    }
}

void LCDWriteNumber(int n, int digits, int x, int y) {
    if (digits < 1) {
        return;
    }
    LCD_set_xy(x, y);
    counts_dispose(n);

    if (digits > 3) {
        lcd_data(thous);
    }
    if (digits > 2) {
        lcd_data(hunds);
    }
    if (digits > 1) {
        lcd_data(tens);
    }
    lcd_data(ones);
}

void LCDWriteChar(char ch, int x, int y) {
    LCD_set_xy(x, y);
    lcd_data(ch);
}

void zxDNMSSwitchSessions(int DNMS, float interOdorDelay, int interti, int trialNumperSess, float waterlen, float delayBeforeReward, int missLimit, int totalSession) {

    DelayNmSecNew(500);

    int currentTrial = 0;
    int currentSession = 0;
    correctRatio = 0;

    while ((currentMiss < missLimit) && (currentSession < totalSession)) {
        if (correctRatio >= 80 && zxAutoStop) {
            home_clr();
            LCDWriteString("Reached 80%");
            while (1);
        } else if (correctRatio >= 80 && zxAutoSwitch) {
            DNMS = DNMS ? 0 : 1;
        }
        currentSession++;
        protectedSerialSend(SpSessN, currentSession);
        protectedSerialSend(SpSess, 1);

        home_clr();
        puts_lcd((unsigned char*) &text7[0], sizeof (text7));
        LCDWriteChar(DNMS ? 'N' : 'M', 1, 1);
        line_2();
        puts_lcd((unsigned char*) &text8[0], sizeof (text8));
        LCDWriteNumber(currentSession, 2, 2, 2);

        currentTrial = 0;
        hit = 0;
        miss = 0;
        falseAlarm = 0;
        correctRejection = 0;
        unsigned int shuffledList[4];
        while (currentTrial < (trialNumperSess) && currentMiss < missLimit) {
            shuffleArray(shuffledList, 4);
            int iter;
            for (iter = 0; iter < 4 && currentMiss < missLimit; iter++) {
                int firstOdor, secondOdor;
                int index = shuffledList[iter];
                firstOdor = (index == 0 || index == 2) ? odor_A : odor_B;
                secondOdor = (index == 1 || index == 2) ? odor_A : odor_B;
                LCD_set_xy(2, 1);
                lcd_data(firstOdor + 0x3F);
                lcd_data(secondOdor + 0x3F);
                zxDNMSSwitchTrial(DNMS, firstOdor, odorLength, interOdorDelay, secondOdor, waterlen, interti, delayBeforeReward);
                currentTrial++;
            }
        }

        protectedSerialSend(SpSess, 0);
    }
    protectedSerialSend(SpTrain, 0);

}

void zxDNMSLaserSessions(int type, float thisdelay, int interti, int trialsPerSession, float WaterLen, int missLimit, int totalSession, float delay_before_reward) {
    protectedSerialSend(SpStepN, 0);
    DelayNmSecNew(2 * 1000);

    int currentTrial = 0;
    int currentSession = 0;

    while ((currentMiss < missLimit) && (currentSession < totalSession)) {
        currentSession++;
        protectedSerialSend(SpOdorDelay, thisdelay);
        protectedSerialSend(SpSess, 1);

        home_clr();
        puts_lcd((unsigned char*) &text7[0], sizeof (text7));
        line_2();
        puts_lcd((unsigned char*) &text8[0], sizeof (text8));

        LCDWriteNumber(currentSession, 2, 2, 2);
        hit = 0;
        miss = 0;
        falseAlarm = 0;
        correctRejection = 0;
        unsigned int shuffledList[4];
        for (currentTrial = 0; currentTrial < trialsPerSession && currentMiss < missLimit;) {
            shuffleArray(shuffledList, 4);
            int iter;
            for (iter = 0; iter < 4 && currentMiss < missLimit; iter++) {
                DelayNmSecNew(1000);
                int firstOdor, secondOdor;
                int index = shuffledList[iter];
                firstOdor = (index == 0 || index == 2) ? odor_A : odor_B;
                secondOdor = (index == 1 || index == 2) ? odor_A : odor_B;
                LCD_set_xy(1, 1);
                lcd_data(firstOdor + 0x3F);
                lcd_data(secondOdor + 0x3F);
                //                int laserOnTrial = currentTrial % 2;
                //                if (ZJAllSessionLaserOn) {
                //                    laserOnTrial = 1;
                //                }
                int laserCurrentTrial;
                switch (laserTrialType) {
                    case laserNoTrial:
                        laserCurrentTrial = 0;
                        break;
                    case laserEveryTrial:
                        laserCurrentTrial = 1;
                        break;
                    case laserOtherTrial:
                        laserCurrentTrial = currentTrial % 2;
                        break;
                    case laserFollowOdorA:
                        laserCurrentTrial = (firstOdor == odor_A);
                        break;
                    case laserFollowOdorB:
                        laserCurrentTrial = (firstOdor == odor_B);
                }
                //                laserCurrentTrial = (laserTrialType == laserNoTrial) ? 0 : (laserTrialType == laserEveryTrial) ? 1 : (currentTrial % 2);
                zxDNMSLaserTrial(type, firstOdor, odorLength, thisdelay, secondOdor, WaterLen, interti, delay_before_reward, laserCurrentTrial);
                currentTrial++;
            }
        }
        protectedSerialSend(SpSess, 0);
    }
    protectedSerialSend(SpTrain, 0); // send it's the end
}

void zxDNMSSwitchTrial(int DNMS, int FirstOdorIn, float odorLength, float OdorDelayIn, int SecondOdorIn, float WaterLIn, int IntervalIn, float delay_before_reward) {
    protectedSerialSend(FirstOdorIn + 7, 1); // odor_A=2, spOdor_A=9
    LCDWriteChar('1', 4, 1);
    Valve_ON(FirstOdorIn, period * 2);
    timerCounterI = 0;
    while (timerCounterI < odorLength * 1000) {
    }
    Valve_OFF(FirstOdorIn);
    protectedSerialSend(FirstOdorIn + 7, 0); // 0 for close
    LCDWriteChar('d', 4, 1);
    timerCounterI = 0;
    while (timerCounterI < OdorDelayIn * 1000) {
    }
    LCDWriteChar('2', 4, 1);
    protectedSerialSend(SecondOdorIn + 7, 1); // 1 for open
    Valve_ON(SecondOdorIn, period * 2);
    timerCounterI = 0;
    while (timerCounterI < odorLength * 1000) {
    }
    Valve_OFF(SecondOdorIn);
    protectedSerialSend(SecondOdorIn + 7, 0); // 0 for close
    LCDWriteChar('D', 4, 1);
    timerCounterI = 0;
    while (timerCounterI < delay_before_reward * 1000) {
    }

    LCDWriteChar('R', 4, 1);
    Flags.lick = 0;
    timerCounterI = 0;
    while (timerCounterI < WaterLIn * 1000) {
        if (LickOnPin) {
            Flags.lick = 1;
            if ((FirstOdorIn != SecondOdorIn && DNMS) || (FirstOdorIn == SecondOdorIn && !DNMS)) {
                Valve_ON(water_sweet, period * 2);
            }
        }
    }
    Valve_OFF(water_sweet);

    if (Flags.lick) {
        Flags.lick = 0;
        if ((FirstOdorIn != SecondOdorIn && DNMS) || (FirstOdorIn == SecondOdorIn && !DNMS)) {
            protectedSerialSend(SpWater_sweet, 0);
            hit++;
            protectedSerialSend(SpHit, 1);
            LCDWriteNumber((int) hit, 3, 6, 1);
        } else {
            falseAlarm++;
            protectedSerialSend(SpFalseAlarm, 1);
            LCDWriteNumber((int) falseAlarm, 3, 6, 2);
        }
    } else {
        if ((FirstOdorIn != SecondOdorIn && DNMS) || (FirstOdorIn == SecondOdorIn && !DNMS)) {
            miss++;
            protectedSerialSend(SpMiss, 1);
            LCDWriteNumber((int) miss, 3, 10, 1);
        } else {
            correctRejection++;
            protectedSerialSend(SpCorrectRejection, 1);
            LCDWriteNumber((int) correctRejection, 3, 10, 2);
        }
    }
    int totalTrials = hit + correctRejection + miss + falseAlarm;
    LCDWriteNumber(totalTrials, 3, 14, 2);
    if (hit + correctRejection > 0) {
        correctRatio = 100 * (hit + correctRejection) / totalTrials;
        LCDWriteNumber(correctRatio, 2, 14, 1);
    }

    LCDWriteChar('I', 4, 1);
    protectedSerialSend(SpITI, 1);
    timerCounterI = 0;
    while (timerCounterI < IntervalIn * 1000) {
    }
    protectedSerialSend(SpITI, 0);
}

void zxDNMSLaserTrial(int type, int FirstOdorIn, float odorLength, float interOdorDelay, int SecondOdorIn, float WaterLIn, int ITI, float delay_before_reward, int laserOnTrial) {
    protectedSerialSend(Splaser, laserOnTrial);
    Valve_ON(FirstOdorIn, period * 2);
    protectedSerialSend(FirstOdorIn + 7, 1);
    LCDWriteChar('1', 4, 1);
    assertLaser(type, atFirstOdorBeginning, laserOnTrial);
    wait_ms(odorLength * 1000);
    Valve_OFF(FirstOdorIn);
    LCDWriteChar('d', 4, 1);
    assertLaser(type, atFirstOdorEnd, laserOnTrial);
    protectedSerialSend(FirstOdorIn + 7, 0);

    ///////////-inter odor interval-/////////////////
    assertLaser(type, atDelayBegin, laserOnTrial);
    wait_ms(1000);
    assertLaser(type, atDelayOneSecIn, laserOnTrial);
    wait_ms(interOdorDelay * 1000 - 2000);
    assertLaser(type, atDelayLastSecBegin, laserOnTrial);
    wait_ms(1000);
    ///////////-Second odor-/////////////////
    assertLaser(type, atSecondOdorBeginning, laserOnTrial);
    Valve_ON(SecondOdorIn, period * 2);
    //    if (SecondOdorIn == odor_A) Out1 = 1;
    //    else if (SecondOdorIn == odor_B) Out2 = 1;
    protectedSerialSend(SecondOdorIn + 7, 1);
    LCDWriteChar('2', 4, 1);
    wait_ms(odorLength * 1000);
    Valve_OFF(SecondOdorIn);
    LCDWriteChar('D', 4, 1);
    assertLaser(type, atSecondOdorEnd, laserOnTrial);

    //    if (SecondOdorIn == odor_A) Out1 = 0;
    //    else if (SecondOdorIn == odor_B) Out2 = 0;

    protectedSerialSend(SecondOdorIn + 7, 0);
    ////////-delay before reward-///////

    wait_ms(delay_before_reward * 1000);
    LCDWriteChar('R', 4, 1);
    if (FirstOdorIn != SecondOdorIn) {
        ///////////-reward-/////////////////
        Flags.lick = 0;

        timerCounterI = 0;
        while (timerCounterI < WaterLIn * 1000) {
            if (LickOnPin) {
                Valve_ON(water_sweet, period * 2);
                Flags.lick = 1;
            }
        }
        Valve_OFF(water_sweet);

        if (Flags.lick) {
            Flags.lick = 0;
            currentMiss = 0;
            hit++;
            protectedSerialSend(SpHit, 1);
            LCDWriteNumber(hit, 3, 6, 1);
        } else {
            currentMiss++;
            miss++;
            protectedSerialSend(SpMiss, 1);
            LCDWriteNumber(miss, 3, 10, 1);
        }
    } else {
        ///////////- punishment-/////////////////
        Flags.lick = 0;
        timerCounterI = 0;
        while (timerCounterI < WaterLIn * 1000) {
            if (LickOnPin) {
                Flags.lick = 1;
            }
        }
        if (Flags.lick) {
            Flags.lick = 0; // Reset flag
            falseAlarm++;
            currentMiss = 0;
            protectedSerialSend(SpFalseAlarm, 1);
            LCDWriteNumber(falseAlarm, 3, 6, 2);
        } else {
            correctRejection++;
            Flags.lick = 0;
            protectedSerialSend(SpCorrectRejection, 1);
            LCDWriteNumber(correctRejection, 3, 10, 2);
        }
    }
    // Total Trials
    int totalTrials = hit + correctRejection + miss + falseAlarm;
    LCDWriteNumber(totalTrials, 3, 14, 2);
    // Discrimination rate
    if (hit + correctRejection > 0) {
        correctRatio = 100 * (hit + correctRejection) / totalTrials;
        LCDWriteNumber(correctRatio, 2, 14, 1);
    }
    LCDWriteChar('I', 4, 1);
    ///--ITI1---///
    assertLaser(type, atITIBeginning, laserOnTrial);
    wait_ms(1000);
    assertLaser(type, atITIOneSecIn, laserOnTrial);
    wait_ms(ITI * 1000 - 1000);
}

void zxGoNogoSessions(int odorType, int autoSwitch, int trialInSession, int delay, float waterlen,                                                                                                                                                                                                                                                                               \
             int ITI, int preOdorLaser, int delaylaseryes) {
    if (preOdorLaser) {
        LCD_set_xy(1, 1);
        LCDWriteString("Laser Freq   Hz");
        int laserFreq = laserOffTime ? 1000 / (laserOnTime + laserOffTime) : 1;
        LCDWriteNumber(laserFreq, 2, 12, 1);
        DelayNmSecNew(1000);
    }

    int go_odor;
    int nogo_odor;
    switch (odorType) {
        case 1:
            go_odor = odor_A;
            nogo_odor = odor_B;
            break;
        case 2:
            go_odor = odor_B;
            nogo_odor = odor_A;
            break;
        case 3:
            go_odor = odor_C;
            nogo_odor = odor_D;
            break;
        case 4:
            go_odor = odor_D;
            nogo_odor = odor_C;
            break;
        default:
            go_odor = odor_A;
            nogo_odor = odor_B;
    }

    DelayNmSecNew(500);

    int sessionNum = 0;
    int lastRatio = 0;
    unsigned int shuffledList[trialInSession];
    while (1) { //while (1)  no stop
        if ((correctRatio >= HighThres) || (correctRatio >= lowThres && lastRatio >= lowThres)) {
            if (autoSwitch == 2) {
                int temp = nogo_odor;
                nogo_odor = go_odor;
                go_odor = temp;
            } else if (autoSwitch == 1) {
                splash("Reach Criteria", "");
                while (1) {
                }
            }
        }
        shuffleArray(shuffledList, trialInSession);
        lastRatio = correctRatio;
        hit = 0;
        miss = 0;
        falseAlarm = 0;
        correctRejection = 0;

        sessionNum++;
        protectedSerialSend(SpSess, 1);
        home_clr();
        puts_lcd((unsigned char*) &text7[0], sizeof (text7));
        line_2();
        puts_lcd((unsigned char*) &text8[0], sizeof (text8));

        LCD_set_xy(1, 1);
        switch (go_odor) {
            case odor_A:
                LCDWriteString("Bgo");
                break;
            case odor_B:
                LCDWriteString("Jgo");
                break;
            case odor_C:
                LCDWriteString("Mgo");
                break;
            case odor_D:
                LCDWriteString("Kgo");
                break;
        }

        LCDWriteNumber(sessionNum, 2, 2, 2);
        int currentTrial;
        int laserOn;
        for (currentTrial = 0; currentTrial < trialInSession; currentTrial++) {
            laserOn = currentTrial % 2;
            int index = shuffledList[currentTrial];
            if (index % 2 == 0) {
                zxGoNogoTrial(go_odor, 1, laserOn, odorLength, delay, waterlen, ITI, preOdorLaser);
            } else {
                zxGoNogoTrial(nogo_odor, 0, laserOn, odorLength, delay, waterlen, ITI, preOdorLaser);
            }
        }
        protectedSerialSend(SpSess, 0);
    }
}

void zxGoNogoTrial(int odor, int isGo, int laserOn, float odorlen, int delay, float waterlen,                                                                                                                                                                                                                                                                               \
             int ITI, int beforeOdorLaser) {

    assertLaser(laserDuringDelay, atDelayOneSecIn, laserOn);
    wait_ms(3000);
    assertLaser(laserDuringDelay, atDelayLastSecBegin, laserOn);
    wait_ms(1000);
    Valve_ON(odor, period * 2);
    protectedSerialSend(odor == odor_A || odor == odor_C ? SpOdor_A : SpOdor_B, 1);
    wait_ms(odorlen * 1000);
    Valve_OFF(odor);
    wait_ms(delay * 1000);


    //---water--//
    Flags.lick = 0;
    timerCounterI = 0;
    while (timerCounterI < waterlen * 1000) {
        if (LickOnPin) {
            if (isGo) {
                Valve_ON(water_sweet, period * 2);
            }
            Flags.lick = 1;
        }
    }

    Valve_OFF(water_sweet);
    if (Flags.lick && isGo) {
        currentMiss = 0;
        hit++;
        protectedSerialSend(SpHit, 1);
        LCDWriteNumber(hit, 3, 6, 1);
    } else if (isGo && !Flags.lick) {
        currentMiss++;
        miss++;
        protectedSerialSend(SpMiss, 1);
        LCDWriteNumber(miss, 3, 10, 1);
    } else if ((!isGo) && Flags.lick) {
        falseAlarm++;
        protectedSerialSend(SpFalseAlarm, 1);
        LCDWriteNumber(falseAlarm, 3, 6, 2);
    } else if ((!isGo) && !Flags.lick) {
        correctRejection++;
        protectedSerialSend(SpCorrectRejection, 1);
        LCDWriteNumber(correctRejection, 3, 10, 2);
    }
    Flags.lick = 0;
    // Total Trials
    int totalTrials = hit + correctRejection + miss + falseAlarm;
    LCDWriteNumber(totalTrials, 3, 14, 2);

    // Discrimination rate
    if (hit + correctRejection > 0) {
        correctRatio = 100 * (hit + correctRejection) / totalTrials;
        LCDWriteNumber(correctRatio, 2, 14, 1);
    }
    //--iti--//

    wait_ms(ITI * 1000);
}

void odorDepeltion(int totalTrial) {
    int n = getFuncNumber(1, "1=BJ 2=MK 3=All");

    int i = 0;
    while (i < totalTrial) {
        home_clr();
        LCDWriteString("Odor Depletion");
        int percent = i > 599 ? 999 : i * 100 / 60;
        LCDWriteNumber(percent, 3, 1, 2);
        LCDWriteString("%, Valve ");
        if (n != 2) {
            LCD_set_xy(15, 2);
            lcd_data(0x31);
            Valve_ON(odor_A, period * 2);
            DelayNmSecNew(15000);
            Valve_OFF(odor_A);
            LCD_set_xy(15, 2);
            lcd_data(0x32);
            Valve_ON(odor_B, period * 2);
            DelayNmSecNew(15000);
            Valve_OFF(odor_B);
        }
        if (n != 1) {
            LCD_set_xy(15, 2);
            lcd_data(0x33);
            Valve_ON(5, period * 2);
            DelayNmSecNew(15000);
            Valve_OFF(5);
            LCD_set_xy(15, 2);
            lcd_data(0x34);
            Valve_ON(6, period * 2);
            DelayNmSecNew(15000);
            Valve_OFF(6);
        }
        i++;
    }
}

void feedWater(int threshold) {
    InitZXTMR();
    int topTime = 600;
    switch (threshold) {
        case 1:
            topTime = 2;
            break;
        case 2:
            topTime = 5;
            break;
        case 3:
            topTime = 10;
            break;
        case 4:
            topTime = 30;
            break;
        case 5:
            topTime = 600;
    }

    splash("Auto Feed Water", "Wait for lick...");
    while (!licking) {
    }
    int go = 1;
    while (go) {
        splash("Feeding", "   s left.");
        Valve_ON(water_sweet, period * 2);
        int secondsWithoutLick = 0;
        int licked = 0;
        while (secondsWithoutLick < topTime) {
            LCDWriteNumber(topTime - secondsWithoutLick, 3, 1, 2);
            timerCounterI = 0;
            while (timerCounterI < 1000) {
                if (LickOnPin) {
                    licked = 1;
                }
            }
            secondsWithoutLick = licked ? 0 : secondsWithoutLick + 1;
            licked = 0;
        }
        Valve_OFF(water_sweet);
        go = (getFuncNumber(1, "No Lick. Again=1") == 1);
    }
}

void feedWaterNew(void) {
    int waterLength = 1000;
    switch (getFuncNumber(1, "1=0.5s 2=1s 3=2s")) {
        case 1:
            waterLength = 500;
            break;
        case 2:
            waterLength = 1000;
            break;
        case 3:
            waterLength = 2000;
    }
    int wasLicking = 0;
    int lickCounter = 0;
    splash("Total Lick", "");
    while (1) {
        if (licking && !wasLicking) {
            timerCounterI = 0;
            Valve_ON(water_sweet, period * 2);
            wasLicking = 1;
            lickCounter++;
            LCDWriteNumber(lickCounter, 4, 12, 1);
        } else if (wasLicking && !licking) {
            wasLicking = 0;
        }
        if (timerCounterI > waterLength) {
            Valve_OFF(water_sweet);
        }
    }
}

void varyDelay(void) {
    int m = getFuncNumber(1, "5s 8s 12s delay");
    float delay = 5.0;
    switch (m) {
        case 1:
            delay = 5.0;
            break;
        case 2:
            delay = 8.0;
            break;
        case 3:
            delay = 12.0;
            break;
        default:
            delay = 5.0;
    }
    laserOnTime = 65535;
    laserOffTime = 0;
    zxDNMSLaserSessions(laserDuringDelay, delay, delay * 2, 20, 0.5, 50, 12, 1.0);
}

void countRounds(void) {
    home_clr();
    int rounds = 0;
    int displayNumber = 0;
    int primed = 0;
    while (1) {
        if (!(PORTDbits.RD14 || PORTDbits.RD15 || primed)) {
            primed = 1;
        } else if (primed && PORTDbits.RD14) {
            rounds++;
            primed = 0;
        } else if (primed && PORTDbits.RD15) {
            rounds--;
            primed = 0;
        }
        displayNumber = rounds > 0 ? rounds : -rounds;

        LCDWriteChar(rounds > 0 ? 'L' : 'R', 1, 1);
        LCDWriteNumber(displayNumber, 4, 2, 1);

        PORTB |= 0x70;
        PORTBbits.RB4 = 0;
        if (PORTBbits.RB0 == 0) {
            rounds = 0;
            primed = 0;
        }
    }
}

void setAutoStop(int autoStop_In) {
    zxAutoStop = autoStop_In;
}

void setAutoSwitch(int autoSwitch_in) {
    zxAutoSwitch = autoSwitch_in;
}

void splash(char * s1, char * s2) {
    home_clr();
    LCDWriteString(s1);
    line_2();
    LCDWriteString(s2);
    DelayNmSecNew(1000);
}

void wait_ms(int time) {
    timerCounterI = 0;
    while (timerCounterI < time) {
    }
}

void variableVoltage() {
    char VString[] = {'1', '2', '.', '0', 'V', ' ', '-', '-', ' ', '-', ' ', '+', ' ', '+', '+'};
    int V_step = 0x85;
    int V = 120;
    while (1) {
        int pwmFreq = V * V_step;
        Valve_ON(4, pwmFreq);
        VString[0] = V / 100 + 0x30;
        VString[1] = ((V % 100) / 10) + 0x30;
        VString[3] = (V % 10) + 0x30;
        int n = getFuncNumber(1, VString);
        switch (n) {
            case 1:
                V -= 5;
                break;
            case 2:
                V -= 1;
                break;
            case 3:
                V += 1;
                break;
            case 4:
                V += 5;
        }
        V = (V > 120) ? 120 : V;
        V = (V < 0) ? 0 : V;
    }
}

void callFunction(int n) {
    InitZXTMR();
    switch (n) {
            int m;
            //ZJ's functions
        case 4111:
            splash("Laser in Delay", "8Hz 25ms On");
            m = getFuncNumber(1, "All OdorA OdorB");
            switch (m) {
                case 1:
                    laserTrialType = laserEveryTrial;
                    break;
                case 2:
                    laserTrialType = laserFollowOdorA;
                    break;
                case 3:
                    laserTrialType = laserFollowOdorB;
                    break;
                default:
                    laserTrialType = laserEveryTrial;
            }
            laserOnTime = 25;
            laserOffTime = 100;

            zxDNMSLaserSessions(laserDuringDelay, 5.0, 10, 20, 0.5, 50, 5, 1.0);
            break;

            //zx's functions
        case 4311:
            n = getFuncNumber(1, "BGo JGo MGo KGo");
            m = getFuncNumber(1, "Stop Switch Go");
            zxGoNogoSessions(n, m, 20, 1, 0.5, 4, 0, 0);
            break;
        case 4312:
            n = getFuncNumber(1, "BGo JGo MGo KGo");
            m = getFuncNumber(1, "Stop Switch Go");
            zxGoNogoSessions(n, m, 20, 1, 0.5, 4, 1, 0);
            break;
        case 4313:
            splash("DNMS Control", "M-Go");
            setLaser();
            zxGoNogoSessions(3, 3, 20, 1, 0.5, 4, 1, 0);
            break;
        case 4321:
            setLaser();
            break;
        case 4322:
            variableVoltage();
            break;
        case 4411:
            n = getFuncNumber(1, "Stop Switch Go");
            setAutoStop(n == 1 ? 1 : 0);
            setAutoSwitch(n == 2 ? 1 : 0);
            n = getFuncNumber(1, "1=DNMS 2=DMS");
            zxDNMSSwitchSessions(n == 1, 5.0, 10, 20, 0.5, 1.0, 100, 40);
            break;
        case 4412:
            splash("During Delay", "");
            setLaser();
            n = getFuncNumber(1, "session 1=12 2=5");
            zxDNMSLaserSessions(laserDuringDelay, 5.0, 10, 20, 0.5, 50, (n == 2 ? 5 : 12), 1.0);
            break;
        case 4413:
            splash("During Puff", "");
            zxDNMSLaserSessions(laserDuringOdor, 5.0, 10, 20, 0.5, 50, 12, 1.0);
            break;
        case 4414:
            splash("During All", "1 ON 1 OFF");
            zxDNMSLaserSessions(laserDuringBeginningToOneSecInITI, 5.0, 10, 20, 0.5, 50, 5, 1.0);
            break;
        case 4415:
            splash("Varying Delay", "DC laser,12 Sess");
            varyDelay();
            break;

        case 4421:
            splash("DNMS Shaping", "No Laser");
            laserTrialType = laserNoTrial;
            zxDNMSShaping();
            break;
        case 4422:
            splash("DNMS Shaping", "1 ON 1 OFF");
            laserTrialType = laserOtherTrial;
            zxDNMSShaping();
            break;
        case 4423:
            splash("DNMSShaping", "Laser On");
            laserTrialType = laserEveryTrial;
            zxDNMSShaping();
            break;
        case 4444:
            feedWaterNew();
            break;
        case 4445:
            countRounds();
            break;


    }
}

